package com.stockprediction.analysis;

import libsvm.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.HashMap;

public class SVMTrainer {
    public static svm_model train(List<double[]> X_train, List<Integer> y_train) {
        svm_problem problem = new svm_problem();
        int dataCount = X_train.size();
        problem.l = dataCount;
        problem.x = new svm_node[dataCount][];
        problem.y = new double[dataCount];

        // 1️⃣ 计算类别权重 (解决数据不均衡问题)
        Map<Integer, Integer> classCounts = new HashMap<>();
        for (int label : y_train) {
            classCounts.put(label, classCounts.getOrDefault(label, 0) + 1);
        }
        double weight0 = 1.0 / classCounts.getOrDefault(0, 1); // 负面
        double weight1 = 1.0 / classCounts.getOrDefault(1, 1); // 正面

        // 2️⃣ 存储非零特征（稀疏优化）
        for (int i = 0; i < dataCount; i++) {
            double[] features = X_train.get(i);
            int nonZeroCount = 0;
            for (double v : features) {
                if (v != 0) nonZeroCount++;
            }

            svm_node[] nodes = new svm_node[nonZeroCount];
            int index = 0;
            for (int j = 0; j < features.length; j++) {
                if (features[j] != 0) {
                    nodes[index] = new svm_node();
                    nodes[index].index = j + 1;
                    nodes[index].value = features[j];
                    index++;
                }
            }

            problem.x[i] = nodes;
            problem.y[i] = y_train.get(i);
        }

        // 3️⃣ 设定 SVM 参数
        svm_parameter param = new svm_parameter();
        param.kernel_type = svm_parameter.LINEAR; // 线性核
        param.C = 1.0;  // 适中 C 值
        param.svm_type = svm_parameter.C_SVC;
        param.probability = 1;
        param.cache_size = 500;  // 增加缓存大小（单位 MB）
        param.eps = 0.001;  // 训练收敛精度
        param.nr_weight = 2; // 设置类别权重数量
        param.weight_label = new int[]{0, 1}; // 负面、正面类别
        param.weight = new double[]{weight0, weight1}; // 计算出的类别权重
        param.kernel_type = svm_parameter.LINEAR;
        param.gamma = 1.0 / X_train.get(0).length;




        svm_model svmModel = svm.svm_train(problem, param);
        try{
            svm.svm_save_model("svm_model.txt",svmModel);
            System.out.printf("模型训练完成，并保存到 svm_model.txt");
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        return svmModel;
    }
}
